#-*-coding: utf-8-*-

# Version: 0.1
# Author: Don Li <iedon@126.com>
# License: Copyright(c) 2017 Don.Li
# Summary:

import os
import multiprocessing
import asyncio
import queue
import sys
import traceback

from log  import log

class BrokenProcessPool(Exception): pass

class Task(object):
    def __init__(self, func, args = None, kwargs = None):
        self.future  = asyncio.Future()
        self.func    = func
        self.args    = args
        self.kwargs  = kwargs

    @property
    def value(self):
        return id(self), self.func, self.args, self.kwargs

class Worker(multiprocessing.Process):
    def __init__(self, q_in, q_out):
        multiprocessing.Process.__init__(self, target=self.__loop, args=(q_in, q_out), daemon=True)
        self.__curr_task_id = multiprocessing.Value('d', 0)

    @property
    def task_id(self):
        return self.__curr_task_id.value

    def __loop(self, q_in, q_out):
        if hasattr(q_in, '_writer'):
            q_in._writer.close()
        if hasattr(q_out, '_reader'):
            q_out._reader.close()

        while 1:
            try:
                task_id, func, args, kwargs = q_in.get()
            except KeyboardInterrupt:
                break

            if task_id >= 0 and func:
                self.__curr_task_id.value = task_id
                try:
                    res = func(*args, **kwargs)
                    throw_excep = 0
                except Exception as e:
                    throw_excep = 1
                    res = traceback.format_exc()

                q_out.put((task_id, res, throw_excep))
                self.__curr_task_id.value = 0

                log.debug("task_id: {0}, func: {1}, args: {2}, kwargs: {3}, res:{4}".format(
                    task_id, func, args, kwargs, res))
            else:
                break

class Pool(object):
    def __init__(self, processes = None):
        self.__loop = asyncio.get_event_loop()

        if not processes:
            processes = os.cpu_count() or 1

        self.__max_processes = processes

        self.__workers  = []
        self.__tasks    = []

        self.__running  = False

    def start(self):
        self.__q_in  = multiprocessing.Queue()
        self.__q_out = multiprocessing.Queue()

        self.__repopulate()

        self.__loop.call_later(0.001, self.__maintain)
        self.__running = True

    def stop(self):
        if self.__running:
            self.__running = False

            for i in reversed(range(len(self.__workers))):
                self.__q_in.put((-1, None, None, None)) # notify worker finish
                worker = self.__workers[i]
                worker.join(1)
                del self.__workers[i]

    def __repopulate(self):
        for i in range(self.__max_processes - len(self.__workers)):
            worker = Worker(self.__q_in, self.__q_out)
            self.__workers.append(worker)
            worker.start()
            log.debug("starting {0} <PID: {1}>".format(worker.name, worker.pid))

    def __join_exited_workers(self):
        cleaned = False

        for i in reversed(range(len(self.__workers))):
            worker = self.__workers[i]

            if worker.exitcode is not None:
                _task_id = worker.task_id
                log.error("Worker<pid:{1},task_id:{2})> was broken! exitcode: {0}, exiting...".format(worker.exitcode, worker.pid, _task_id))
                if _task_id:
                    _task = self.task(_task_id, True)
                    if _task:
                        _task.future.set_exception(BrokenProcessPool())

                worker.join()
                del self.__workers[i]
                cleaned = True

        return cleaned

    def apply(self, func, *args, **kwargs):
        timeout = kwargs.pop('timeout', None)

        task = Task(func, args, kwargs)
        self.__tasks.append(task)
        self.__q_in.put(task.value)

        return asyncio.wait_for(task.future, timeout = timeout)

    def __maintain(self):
        self.__recv()

        if self.__running:
            if self.__join_exited_workers():
                self.__repopulate()

            self.__loop.call_later(0.001, self.__maintain)

    def task(self, task_id, pop = False):
        for idx, task in enumerate(self.__tasks):
            if task_id == id(task):
                if pop:
                    del self.__tasks[idx]

                return task
        else:
            return None

    def __recv(self):
        while 1:
            try:
                task_id, res, throw_excep = self.__q_out.get_nowait()
                task = self.task(task_id, True)
                if task and task.future:
                    if task.future.done():
                        log.warn("Task<id:{0}, done:{3}> timeout. res:{1}, throw_excep: {2}".format(
                            task_id, res, throw_excep, task.future.done()))
                    else:
                        if throw_excep:
                            task.future.set_exception(Exception(res))
                        else:
                            task.future.set_result(res)
                else:
                    log.debug("missed task:{0}, res:{1}, throw_excep: {2}".format(task_id, res, throw_excep))
            except:
                break
